from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Activation
import tensorflow

def create_model():
    MNISTModel = Sequential()
    MNISTModel.add(Dense(128, input_shape=(784,)))
    MNISTModel.add(Activation('relu'))
    MNISTModel.add(Dense(10))

    MNISTModel.add(Activation('softmax'))
    # lr为学习率，epsilon防止出现0，rho/decay分别对应公式中的beta_1和beta_2
    rmsprop = RMSprop(lr=0.001, rho=0.9, epsilon=1e-08, decay=0.00001)
    # 编译模型 主要是确定优化方法，损失函数等
    MNISTModel.compile(optimizer=rmsprop, loss='categorical_crossentropy', metrics=['accuracy'])
    # 自动打印出模型概述信息
    # MNISTModel.summary()
    return MNISTModel
